package utils;

import inference.TableInferer;

import java.text.DecimalFormat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import model.ModelParameters;
import model.TableInformation;

import com.analog.lyric.dimple.factorfunctions.core.IFactorTable;
import com.analog.lyric.dimple.model.Discrete;

public class PrettyPrinter {
	private static DecimalFormat PROB_FORMAT = new DecimalFormat("0.000");
	
	public static void prettyPrintTableInterpretation(TableInferer infer, int[] backbone, double[][][] entity_vars_beliefs,
			double[][] predicate_vars_beliefs, Discrete topic_var, double[] topic_beliefs, 
			Discrete topic_type_var, double[] topic_type_beliefs, int max_printed) {
		ModelParameters model_params = infer.getModelParameters();
		TableInformation table_info = model_params.getTableInformation();
		int backbone_size = ArrayUtils.indexOf(backbone, -2);
		if (backbone_size == -1) {
			backbone_size = backbone.length;
		}
		System.out.println(backbone_size);
		String[][][] sorted_output_entity_domains = new String[table_info.numRows()][backbone_size][];
		double[][][] sorted_output_entity_probs = new double[table_info.numRows()][backbone_size][];
		String[][] sorted_output_predicate_domains = new String[backbone_size][];
		double[][] sorted_output_predicate_probs = new double[backbone_size][];
		for (int col = 0; col < backbone_size; col++) {
			for (int row = 0; row < table_info.numRows(); row++) {
				sorted_output_entity_domains[row][col] = ArrayUtils.translateIdArrayToStringArray(
						model_params.getEntityVariable(row, col).getDomain().getElements(),
						table_info);

				Pair<String[], double[]> sorted_pair = ArrayUtils.sortPairedArrays(
						sorted_output_entity_domains[row][col],
						entity_vars_beliefs[row][col]);
				sorted_output_entity_domains[row][col] = sorted_pair.getFirst();
				sorted_output_entity_probs[row][col] = sorted_pair.getSecond();
			}
			sorted_output_predicate_domains[col] = ArrayUtils.translateIdArrayToStringArray(
					model_params.getPredicateVariable(col).getDomain().getElements(),
					table_info);
			Pair<String[], double[]> sorted_pair = ArrayUtils.sortPairedArrays(
					sorted_output_predicate_domains[col],
					predicate_vars_beliefs[col]);
			sorted_output_predicate_domains[col] = sorted_pair.getFirst();
			sorted_output_predicate_probs[col] = sorted_pair.getSecond();
		}
		
		String[] sorted_output_topic_domain = ArrayUtils.translateIdArrayToStringArray(topic_var.getDomain().getElements(),
				table_info);
		Pair<String[], double[]> sorted_pair = ArrayUtils.sortPairedArrays(sorted_output_topic_domain, topic_beliefs);
		sorted_output_topic_domain = sorted_pair.getFirst();
		double[] sorted_output_topic_probs = sorted_pair.getSecond();
		
		String[] sorted_output_topic_type_domain = ArrayUtils.translateIdArrayToStringArray(topic_type_var.getDomain().getElements(),
				table_info);
		sorted_pair = ArrayUtils.sortPairedArrays(sorted_output_topic_type_domain, topic_type_beliefs);
		sorted_output_topic_type_domain = sorted_pair.getFirst();
		double[] sorted_output_topic_type_probs = sorted_pair.getSecond();
		
		
		int[] max_col_width = new int[backbone_size];
		int[] max_row_height = new int[table_info.numRows()];
		for (int row = 0; row < table_info.numRows(); row++) {
			max_row_height[row] = 0;
		}
		int max_num_preds = 0;
		
		for (int col = 0; col < backbone_size; col++) {
			String[] vals = sorted_output_predicate_domains[col];
			if (vals.length > max_num_preds) {
				max_num_preds = vals.length;
			}
			int max_pred_length = -1;
			for (int val = 0; val < vals.length; val++) {
				int string_len = vals[val].length();
				if (string_len > max_pred_length) {
					max_pred_length = string_len;
				}
			}
			int max_entity_length = -1;
			for (int row = 0; row < table_info.numRows(); row++) {
				vals = sorted_output_entity_domains[row][col];
				if (vals.length > max_row_height[row]) {
					if (vals.length > max_printed)
						max_row_height[row] = max_printed;
					else
						max_row_height[row] = vals.length;
				}
				for (int val = 0; val < vals.length; val++) {
					int string_len = vals[val].length();
					if (string_len > max_entity_length) {
						max_entity_length = string_len;
					}
				}
			}
			max_col_width[col] = Math.max(max_pred_length, max_entity_length);
		}
		int table_width = ((table_info.numVisibleCols() + 1) * 7) + 1;
		for (int col = 0; col < backbone_size; col++) {
			table_width += max_col_width[col];
		}
		
		if (model_params.ignoreTableTopic()) {
		} else {
			for (int topic = 0; topic < sorted_output_topic_domain.length; topic++) {
				System.out.println(PROB_FORMAT.format(sorted_output_topic_probs[topic]) + "  " + sorted_output_topic_domain[topic]);
			}
			
			for (int type = 0; type < sorted_output_topic_type_domain.length; type++) {
				System.out.println(PROB_FORMAT.format(sorted_output_topic_type_probs[type]) + "  " + sorted_output_topic_type_domain[type]);
			}
			
			// Pretty print the column probabilities.
			for (int i = 0; i < max_num_preds; i++) {
				printColumnProb(backbone_size, sorted_output_predicate_domains,
						sorted_output_predicate_probs, i, max_col_width, table_info);
//				if (true)
//					return;
			}
			
			printStringNTimes("=", table_width);
			System.out.println();
			
			for (int row = 0; row < table_info.numRows(); row++) {
				for (int i = 0; i < max_row_height[row]; i++) {
					printRowProb(sorted_output_entity_domains[row], sorted_output_entity_probs[row],
							i, backbone_size,
							max_col_width, table_info);
				}
				System.out.print("|");
				printStringNTimes("-", (table_width - 2));
				System.out.println("|");
			}
			
			System.out.println();
		}
	}
	
	private static void printRowProb(String[][] sorted_output_entity_domains, double[][] sorted_output_entity_probs,
			int i, int num_cols, int[] max_col_width, TableInformation table_info) {
		System.out.print("|");
		for (int col = 0; col < num_cols; col++) {
			if (sorted_output_entity_domains[col].length > i) {
				String val = sorted_output_entity_domains[col][i];
				double belief = sorted_output_entity_probs[col][i];
				if (belief == 0.0) {
					printStringNTimes(" ", (max_col_width[col] + 6));
				} else {
					int num_spaces = max_col_width[col] - val.length();
					printStringNTimes(" ", num_spaces);
					System.out.print(val);
					System.out.print("=");
					System.out.print(PROB_FORMAT.format(belief));
				}
			} else {
				printStringNTimes(" ", (max_col_width[col] + 6));
			}
			System.out.print("|");
		}
		System.out.println();
	}
	
	private static void printColumnProb(int num_cols,
			String[][] sorted_output_pred_domain, double[][] sorted_output_pred_belief,
			int i, int[] max_col_width, TableInformation table_info) {
		System.out.print("|");
		for (int col = 0; col < num_cols; col++) {
			if (sorted_output_pred_domain[col].length > i) {
				String val = sorted_output_pred_domain[col][i];
				double belief = sorted_output_pred_belief[col][i];
				if (belief == 0.0) {
					printStringNTimes(" ", (max_col_width[col] + 6));
				} else {
					int num_spaces = max_col_width[col] - val.length();
					printStringNTimes(" ", num_spaces);
					System.out.print(val);
					System.out.print("=");
					System.out.print(PROB_FORMAT.format(belief));
				}
			} else {
				printStringNTimes(" ", (max_col_width[col] + 6));
			}
			System.out.print("|");
		}
		System.out.println();
	}
	
	private static void printStringNTimes(String str, int n) {
		for (int i = 0; i < n; i++) {
			System.out.print(str);
		}
	}
	
	public static void prettyPrintPredicateFactor(IFactorTable table, Discrete type_var, Discrete pred_var,
			TableInformation table_info) {
		int[][] indices = table.getIndicesSparseUnsafe();
		double[] probs = table.getWeightsSparseUnsafe();
		Object[] type_elems = type_var.getDomain().getElements();
		Object[] pred_elems = pred_var.getDomain().getElements();
		
		HashMap<Integer, Double> cnts = new HashMap<Integer, Double>();
		
		for (int i = 0; i < indices.length; i++) {
			int[] line = indices[i];
			Double cnt = cnts.get((Integer)(type_elems[line[0]]));
			if (cnt == null)
				cnts.put((Integer)(type_elems[line[0]]), probs[i]);
			else
				cnts.put((Integer)(type_elems[line[0]]), probs[i] + cnt.doubleValue());
		}
		
		for (int i = 0; i < indices.length; i++) {
			int[] line = indices[i];
			Double cnt = cnts.get((Integer)(type_elems[line[0]]));
//			if (Math.abs(cnt.doubleValue() - 1.0) > 0.000001) {
//			if (table_info.getStringForId((Integer)(type_elems[line[0]])).equals("/award/award_honor")) {
				System.out.print(table_info.getStringForId((Integer)(type_elems[line[0]])) + "  " 
				+ table_info.getStringForId((Integer)(pred_elems[line[1]])));
				System.out.println("  " + probs[i]);
//			}
		}
	}
	
	public static void prettyPrintPredicateFactor(int[][] indices, double[] probs, Discrete type_var, Discrete pred_var,
			TableInformation table_info) {
		Object[] type_elems = type_var.getDomain().getElements();
		Object[] pred_elems = pred_var.getDomain().getElements();
		
		HashMap<Integer, Double> cnts = new HashMap<Integer, Double>();
		
		for (int i = 0; i < indices.length; i++) {
			int[] line = indices[i];
			Double cnt = cnts.get((Integer)(type_elems[line[0]]));
			if (cnt == null)
				cnts.put((Integer)(type_elems[line[0]]), probs[i]);
			else
				cnts.put((Integer)(type_elems[line[0]]), probs[i] + cnt.doubleValue());
		}
		
		for (int i = 0; i < indices.length; i++) {
			int[] line = indices[i];
			Double cnt = cnts.get((Integer)(type_elems[line[0]]));
//			if (Math.abs(cnt.doubleValue() - 1.0) > 0.000001) {
				System.out.print(table_info.getStringForId((Integer)(type_elems[line[0]])) + "  " 
				+ table_info.getStringForId((Integer)(pred_elems[line[1]])));
				System.out.println("  " + probs[i]);
//			}
		}
	}
	
	public static void prettyPrintEntityFactor(IFactorTable table, Discrete subj_var, Discrete pred_var,
			Discrete obj_var, TableInformation table_info) {
		int[][] indices = table.getIndicesSparseUnsafe();
		double[] probs = table.getWeightsSparseUnsafe();
		Object[] subj_elems = subj_var.getDomain().getElements();
		Object[] pred_elems = pred_var.getDomain().getElements();
		Object[] obj_elems = obj_var.getDomain().getElements();
		
		HashMap<IntPair, Double> cnts = new HashMap<IntPair, Double>();
		for (int i = 0; i < indices.length; i++) {
			int[] line = indices[i];
//			System.out.println(line[0] + " " + line[1] + " " + line[2]);
			IntPair new_pair = new IntPair((Integer)(subj_elems[line[0]]), (Integer)(pred_elems[line[1]]));
			Double cnt = cnts.get(new_pair);
			if (cnt == null) {
				cnts.put(new_pair, probs[i]);
			} else {
				cnts.put(new_pair, (probs[i] + cnt.doubleValue()));
			}
		}

		
		for (int i = 0; i < indices.length; i++) {
			int[] line = indices[i];
			IntPair new_pair = new IntPair((Integer)(subj_elems[line[0]]), (Integer)(pred_elems[line[1]]));
			Double cnt = cnts.get(new_pair);
//			if (Math.abs(cnt - 1.0) > 0.0000000001) {
				System.out.print(table_info.getStringForId((Integer)(subj_elems[line[0]])) + " " + table_info.getStringForId((Integer)(pred_elems[line[1]]))
						+ " " + table_info.getStringForId((Integer)(obj_elems[line[2]])));
				System.out.println("  " + probs[i]);
//			}
		}
	}
}
